package com.gateway.app.lb;

import com.alibaba.fastjson2.JSON;
import com.gateway.app.common.RspEntity;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.time.LocalDate;
import java.time.ZoneId;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 按天 限流控制
 */
@Slf4j
public class DayRequestRateLimiterGatewayFilterFactory extends AbstractGatewayFilterFactory<DayRequestRateLimiterGatewayFilterFactory.Config> {

    private final Map<Long, Map<String, AtomicInteger>> rateLimitMap = new ConcurrentHashMap<>();

    public DayRequestRateLimiterGatewayFilterFactory() {
        super(Config.class);
    }

    @Override
    public List<String> shortcutFieldOrder() {
        // yml配置文件中参数的赋值顺序
        return Arrays.asList("limit");
    }

    @Override
    public GatewayFilter apply(Config config) {
        return (exchange, chain) -> {
            String path = exchange.getRequest().getPath().toString();
            long nowSecond = LocalDate.now().atStartOfDay().atZone(ZoneId.systemDefault()).toInstant().getEpochSecond();
            log.info("nowSecond:{}",nowSecond);
            Map<String, AtomicInteger> apiCount = rateLimitMap.computeIfAbsent(nowSecond, k -> new HashMap<String, AtomicInteger>() {{
                put(path, new AtomicInteger(1));
            }});
            AtomicInteger count = apiCount.computeIfAbsent(path, k -> new AtomicInteger(1));
            int countValue = count.get();
            log.info("count: {}", countValue);
            if (countValue > config.getLimit()) {

                log.warn("已限流: {}", path);
                ServerHttpResponse httpResponse = exchange.getResponse();
                //设置 返回状态码
                httpResponse.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                if (!httpResponse.getHeaders().containsKey("Content-Type")) {
                    httpResponse.getHeaders().add("Content-Type", "application/json");
                }
                RspEntity easyMsResponseEntity = RspEntity.ofFail("0009", "request has been limited. Please try again later");
                //easyMsResponseEntity.setTraceId(EasyMsTraceSynchronizationManager.getTraceId());
                easyMsResponseEntity.setErrorDesc(easyMsResponseEntity.getRetMsg());
                String limitedError = JSON.toJSONString(easyMsResponseEntity);
                log.error("path limitedError error:{}",limitedError);
                DataBuffer buffer = httpResponse.bufferFactory().wrap(limitedError.getBytes(StandardCharsets.UTF_8));
                return httpResponse.writeWith(Mono.just(buffer));
                //exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                //exchange.getResponse().setComplete();
            }
            count.addAndGet(1);
            return chain.filter(exchange).then(Mono.fromRunnable(() -> {
                //删除 过期的key
                rateLimitMap.keySet().stream().filter(time -> time < nowSecond).forEach(rateLimitMap::remove);
            }));
        };
    }

    @Data
    public static class Config {
        private Integer limit;
    }


}
